from collections.abc import Sequence

from pydantic import BaseModel

from onyx.chat.models import LlmDoc
from onyx.context.search.models import InferenceChunk


class DocumentIdOrderMapping(BaseModel):
    order_mapping: dict[str, int]


def map_document_id_order(
    chunks: Sequence[InferenceChunk | LlmDoc], one_indexed: bool = True
) -> DocumentIdOrderMapping:
    order_mapping = {}
    current = 1 if one_indexed else 0
    for chunk in chunks:
        if chunk.document_id not in order_mapping:
            order_mapping[chunk.document_id] = current
            current += 1

    return DocumentIdOrderMapping(order_mapping=order_mapping)


def map_document_id_order_v2(fetched_docs: list[LlmDoc]) -> DocumentIdOrderMapping:
    order_mapping = {}
    for doc in fetched_docs:
        if doc.document_id not in order_mapping and doc.document_citation_number:
            order_mapping[doc.document_id] = doc.document_citation_number
    return DocumentIdOrderMapping(order_mapping=order_mapping)
